import cv2
import numpy as np
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
import re

# 读取图片
img = cv2.imread('screenshot.png')
if img is None:
    raise FileNotFoundError('screenshot.png not found.')

# 获取图像尺寸
height, width = img.shape[:2]
PIECESCOUNT = 3
# 水平和垂直方向各分割成3份
rows = np.array_split(img, PIECESCOUNT, axis=0)
pieces = []
for row in rows:
    cols = np.array_split(row, PIECESCOUNT, axis=1)
    pieces.extend(cols)

# 保存16份小图像（可选）
# 创建线程锁以保证线程安全
lock = threading.Lock()


# 将图像块处理逻辑提取为独立函数
def process_piece(i, piece):
    success, img_encoded = cv2.imencode(".png", piece)
    if success:
        img_bytes = img_encoded.tobytes()
        # 发送POST请求
        response = requests.post("https://home.u96242.nyat.app:14861/ocr",
                                 files={"imgbin": img_bytes},
                                 verify=False)
        try:
            respjson = response.json()
            print(f"Response for piece {i}: {respjson}")
            text_list = respjson[0]
            rect_list = respjson[1]

            # # 使用线程锁保护图像操作
            # with lock:
            #     for j, rect in enumerate(rect_list):
            #         points = np.array(rect, dtype=np.int32)
            # cv2.polylines(piece, [points],
            #               isClosed=True,
            #               color=(0, 255, 0),
            #               thickness=2)
            # # 在边界框附近添加文字
            # if len(points) > 0:
            #     # 选择边界框的左上角点作为文字起始点
            #     text_x, text_y = points[0]
            #     cv2.putText(piece, text_list[j], (text_x, text_y),
            #                 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0),
            #                 1)

            # 保存带边界框的图像
            # cv2.imwrite(f'piece_with_boxes_{i}.png', piece)

            return i, piece, text_list, rect_list
        except ValueError:
            print(f"Response for piece {i} is not valid JSON")
            return i, piece, [], []
    return i, piece, [], []


# 使用多线程处理所有图像块
with ThreadPoolExecutor(max_workers=4) as executor:
    # 提交所有任务
    future_to_piece = {
        executor.submit(process_piece, i, piece): (i, piece)
        for i, piece in enumerate(pieces)
    }

    # 收集结果并处理
    results = []
    for future in as_completed(future_to_piece):
        result = future.result()
        if result:
            results.append(result)

    # 按照原始顺序处理结果
    results.sort(key=lambda x: x[0])

    # 计算原图总面积
    total_area = height * width
    # 用于累计所有标记框的面积
    total_bounding_area = 0

    # 在原图上绘制所有边界框
    for i, piece, text_list, rect_list in results:
        # 计算当前piece在原图中的位置
        row_index = i // PIECESCOUNT
        col_index = i % PIECESCOUNT
        piece_height, piece_width = piece.shape[:2]

        # 将边界框坐标转换为原图坐标
        with lock:  # 使用线程锁保护对原图的操作
            for j, rect in enumerate(rect_list):

                if re.findall(r'下.*河岸', text_list[j]):
                    points = np.array(rect, dtype=np.int32)
                    # 添加偏移量
                    points[:, 0] += col_index * piece_width
                    points[:, 1] += row_index * piece_height
                    cv2.polylines(img, [points],
                                  isClosed=True,
                                  color=(0, 255, 0),
                                  thickness=2)
                    print(text_list[j])
                    # 在原图上绘制对应的文字
                    if len(points) > 0:
                        text_x, text_y = points[0]
                        text_x += col_index * piece_width
                        text_y += row_index * piece_height
                        cv2.putText(img, text_list[j], (text_x, text_y),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0),
                                    1)

                    # 计算当前边界框的面积并累加
                    bounding_area = cv2.contourArea(points)
                    total_bounding_area += bounding_area
                    # 输出当前边界框的面积信息
                    print(
                        f"Text: {text_list[j]}, Bounding box area: {bounding_area}"
                    )
                    
                    # 在满足条件的文字处绘制一个320x240的方框，文字位于右下角
                    # 获取文字的尺寸
                    (text_width, text_height), baseline = cv2.getTextSize(
                        text_list[j], cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                    
                    # 计算320x240方框的左上角坐标，使文字位于右下角
                    rect_x = points[0][0] - 320  # 方框左上角x坐标
                    rect_y = points[0][1] - 240  # 方框左上角y坐标
                    
                    # 绘制320x240的矩形框
                    cv2.rectangle(img, (rect_x, rect_y), 
                                  (rect_x + 320, rect_y + 240), 
                                  (0, 0, 255), 2)
                    
                    # 在方框的右下角绘制文字
                    text_x = rect_x + 320 - text_width - 5  # 右下角x坐标，留5像素边距
                    text_y = rect_y + 240 - 5  # 右下角y坐标，留5像素边距
                    
                    cv2.putText(img, text_list[j], (text_x, text_y),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

    # 计算并输出标记框面积占原图的百分比
    percentage = (total_bounding_area / total_area) * 100
    print(f"Total bounding box area: {total_bounding_area}")
    print(f"Total image area: {total_area}")
    print(f"Bounding boxes area percentage: {percentage:.2f}%")

# 保存带有所有边界框的原图
cv2.imwrite('full_image_with_boxes.png', img)
# Bounding boxes area percentage: 0.38%
# Bounding boxes area percentage: 0.14%
# Bounding boxes area percentage: 0.06%